    clear

% All joints are revolute in this version
% The first body is attached to ground, the last body is free
% Note: All lines of code marked with XS correspond to calculations for the
% second handle of this system. Since this system does not have an
% algebraic constraint on the terminal body, all the calculations
% corresponding to the second handle could have been avoided. All lines of
% codes marked with XS can be commented out and the code will still work
% and give the same value of udot.

% Fixed parameters for the system. Each body in the system is identical
NBODIES = 16;
MASS = 1;
INERTIA = 1/12;
Tor = (1:NBODIES)';
T = 0*Tor;

% For a revolute joint about positive z axis
% H is referred to as P^J_k in our paper. 
H = [0 0 1 0 0 0]';

% D = D^J_k
D = [1 0 0 0 0;
     0 1 0 0 0;
     0 0 0 0 0;
     0 0 1 0 0;
     0 0 0 1 0;
     0 0 0 0 1];

% Center of mass located at the center of each link.
LENGTH = 1;

GRAVITY = [0 -9.81 0]';

% Initialize the state vector of the system
q(1:NBODIES,1) = 1;
u(1:NBODIES,1) = 2;

% PREPROCESSING loop. Operations in this loop are performed just once and
% need not be timed.
for i=1:NBODIES
    
    % Initialize the fixed parameters of the system
    mass(i) = MASS;
    inertia(:,:,i) = [1 0 0;
                      0 1 0;
                      0 0 INERTIA];
                  
    % Spatial mass matrix written in body basis about the center of mass
    M(:,:,i) = [inertia(:,:,i) zeros(3,3);
                zeros(3,3) mass(i)*eye(3,3)];
    
    length(i) = LENGTH;
    
    % Vectors locating handles 1 and 2 with respect to center of mass
    % Written in body basis. 
    cmtoH1(:,i) = [-length(i)/2 0 0]';
    cmtoH2(:,i) = [length(i)/2 0 0]';
    
    % Written in body basis
    cmtoH1cross(:,:,i) = crossmat(cmtoH1(:,i)); 
                     
    cmtoH2cross(:,:,i) = crossmat(cmtoH2(:,i));
                     
    % Shift matrices for spatial constraint forces on handles H1 and H2
    SfH1(:,:,i) = [eye(3) cmtoH1cross(:,:,i);zeros(3) eye(3)];
    SfH2(:,:,i) = [eye(3) cmtoH2cross(:,:,i);zeros(3) eye(3)];
    
    % handle equations for each body are 
    % A_1 = zeta11 F1 + zeta12 F2 + zeta13
    % A_2 = zeta21 F1 + zeta22 F2 + zeta23
    % phi11, phi12, phi21 and phi22 are constant in body basis and computed
    % just once as a preprocessing step. 
    
    shift_inverse_H1(:,:,i) = (SfH1(:,:,i)')*inv(M(:,:,i));
    shift_inverse_H2(:,:,i) = (SfH2(:,:,i)')*inv(M(:,:,i));
    
    zeta11(:,:,i) = shift_inverse_H1(:,:,i)*SfH1(:,:,i); % (1a)
    zeta12(:,:,i) = shift_inverse_H1(:,:,i)*SfH2(:,:,i); % (1a)
    
    zeta21(:,:,i) = shift_inverse_H2(:,:,i)*SfH1(:,:,i); % (1b)
    zeta22(:,:,i) = shift_inverse_H2(:,:,i)*SfH2(:,:,i); % (1b)      
end


for i=1:NBODIES
    % Transformation matrix between body and parent
    pr_c_k(:,:,i) = [cos(q(i)) -sin(q(i)) 0;
                     sin(q(i)) cos(q(i)) 0;
                     0         0         1];
                 
    % Form spatial trasnformation matrices
    Pr_C_K(:,:,i) = [pr_c_k(:,:,i) zeros(3,3);
                    zeros(3,3) pr_c_k(:,:,i)];
    
    % Transformation matrix between body and newtonian frame
    n_c_k(:,:,i) = [cos(sum(q(1:i))) -sin(sum(q(1:i))) 0;
                     sin(sum(q(1:i))) cos(sum(q(1:i))) 0;
                     0                0                1];
   N_C_K(:,:,i) = [n_c_k(:,:,i) zeros(3,3);
                   zeros(3,3) n_c_k(:,:,i)];
               
    % Angular velocity of each body is trivially obtained here
    W(:,i) = [0 0 sum(u(1:i))]';
    
    Fa  = [ -cross(W(:,i),inertia(:,:,i)*W(:,i)); % -w x Iw term
                  mass(i)*n_c_k(:,:,i)'*GRAVITY];% Only external force is gravity
              
    At_H1  = [zeros(3,1);cross(W(:,i),cross(W(:,i),cmtoH1(:,i)))];
    At_H2  = [zeros(3,1);cross(W(:,i),cross(W(:,i),cmtoH2(:,i)))];
    
    if(i~=NBODIES)
        F_Tilde = Fa + H*(T(i,1) -T(i+1,1));
    else
        F_Tilde = Fa + H*T(i,1);
    end
    
    % phi13 and phi23 are also written in body basis
    zeta13(:,i) = shift_inverse_H1(:,:,i)*F_Tilde + At_H1; % (1a)
    zeta23(:,i) = shift_inverse_H2(:,:,i)*F_Tilde + At_H2; % (1b)
    
    % At this point handle equations of motion are ready for all the
    % bodies
end

% A convenient tmp variable.
zeta11_asm_tmp = zeta11(:,:,NBODIES);
zeta12_asm_tmp = zeta12(:,:,NBODIES);
zeta13_asm_tmp = zeta13(:,NBODIES);

zeta21_asm_tmp = zeta21(:,:,NBODIES); % XS
zeta22_asm_tmp = zeta22(:,:,NBODIES); % XS
zeta23_asm_tmp = zeta23(:,NBODIES);   % XS
 
% Assembly
% NBODIES are connected by (NBODIES-1) joint excluding the one which connects
% the first body to ground (if at all)
for i=NBODIES:-1:2   
        
    % Convert handle equations of body i to basis of (i-1)
    % Handle 1
    zeta11_c = Pr_C_K(:,:,i)*zeta11_asm_tmp*Pr_C_K(:,:,i)';
    zeta12_c = Pr_C_K(:,:,i)*zeta12_asm_tmp*Pr_C_K(:,:,i)';
    zeta13_c = Pr_C_K(:,:,i)*zeta13_asm_tmp;
    
    % Handle 2
    zeta21_c = Pr_C_K(:,:,i)*zeta21_asm_tmp*Pr_C_K(:,:,i)';% XS
    zeta22_c = Pr_C_K(:,:,i)*zeta22_asm_tmp*Pr_C_K(:,:,i)';% XS
    zeta23_c = Pr_C_K(:,:,i)*zeta23_asm_tmp;               % XS
    
    % Assemble bodies i and (i-1)   
    X_tilde = inv(D'*(zeta11_c + zeta22(:,:,i-1))*D); % (3e)
    X = D*X_tilde*D'; % (3d)
    
    Wa(:,:,i) = X*zeta21(:,:,i-1);        % (3b)
    Wb(:,:,i) = -X*zeta12_c;               % (3b)
    Wc(:,i) = X*(zeta23(:,i-1)-zeta13_c); % (3c)    
    % Note in equation (3c), \dot{P}^J_K u is identically equal to zero for
    % all planar revolute joints and not included in calculations.
    % PS: The sign of equation Wb in (3b) in the paper draft i sent you  
    % before was wrong.
    
    % Below are equations (4a)-(4f)
    zeta11_asm(:,:,i) = zeta11(:,:,i-1) - zeta12(:,:,i-1)*Wa(:,:,i);
    zeta12_asm(:,:,i) = -zeta12(:,:,i-1)*Wb(:,:,i);
    zeta13_asm(:,i) = zeta13(:,i-1) - zeta12(:,:,i-1)*Wc(:,i);
    
    zeta21_asm(:,:,i) = zeta21_c*Wa(:,:,i);             % XS
    zeta22_asm(:,:,i) = zeta22_c + zeta21_c*Wb(:,:,i);  % XS
    zeta23_asm(:,i) = zeta23_c + zeta21_c*Wc(:,i);      % XS
    
    % Note, the above is a fairly quick and dirty implementatioin and
    % zeta_ij_asm(:,:,1) along with Wa(:,:,1),Wb(:,:,1),Wc(:,1) are not
    % really used. Can be much cleaner by prealloating space for a c++
    % objects containing the _asm variable along with Wa, Wb and Wc. All of
    % these are written in the basis of the parent (or index k) of hte
    % assembly between bodies k and k+1
    
    zeta11_asm_tmp = zeta11_asm(:,:,i);
    zeta12_asm_tmp = zeta12_asm(:,:,i);
    zeta13_asm_tmp = zeta13_asm(:,i);
    
    zeta21_asm_tmp = zeta21_asm(:,:,i); % XS
    zeta22_asm_tmp = zeta22_asm(:,:,i); % XS
    zeta23_asm_tmp = zeta23_asm(:,i);   % XS
    
end

% Spatial acceleration of ground/parent
Ah2_parent = zeros(6,1);

% Disassembly
for i=1:NBODIES
    
    Ah2_parent_inchildbasis = Pr_C_K(:,:,i)'*Ah2_parent;
    
    if(i~=NBODIES)
        % For all intermediate assembled systems, the spatial constraint
        % force on handle 2 is zero. Hence the equation of handle 1 of an
        % assembled intermediate system can be written as 
        % Ah1 = zeta11_asm*Fh1 + zeta13
        % Since D'*(Ah1 - Ah2_parent) = 0 the above equation reduces to
        % D'*(zeta11_asm*Fh1 + zeta13-Ah2_parent)=0 From this, Fh1 can be
        % calculated as 
        Fh1 = D*((D'*zeta11_asm(:,:,i+1)*D) \ (-D'*(zeta13_asm(:,i+1)-Ah2_parent_inchildbasis)));
        
        % the stored values of Wa,Wb and Wc are used to calculate the
        % constrained force on hte connecting joint.
        Fh2 = -Wa(:,:,i+1)*Fh1  -Wc(:,i+1); % This is the negative of eq (3a)        
        Ah1 = zeta11_asm(:,:,i+1)*Fh1 + zeta13_asm(:,i+1);
    else
        % For the last body, there is no assembly.
        Fh1 = D*((D'*zeta11(:,:,i)*D) \ (-D'*(zeta13(:,i)- Ah2_parent_inchildbasis)));
        Ah1 = zeta11(:,:,i)*Fh1 + zeta13(:,i);
        Fh2 = zeros(6,1);
    end
    
    Ah2 = zeta21(:,:,i)*Fh1 + zeta22(:,:,i)*Fh2 + zeta23(:,i);
    
    % w_parent x w_child terms is identically equal to zero and not taken
    % into account here
    udot(i,1) = H'*(Ah1 - Ah2_parent_inchildbasis);    
    
    Ah2_parent = Ah2;
end

DD = [0 0 1 0 0 0 ;
     0 0 0 1 0 0 ;
     0 0 0 0 1 0]';